import os
import shutil
import random
from tqdm import tqdm
import numpy as np

def set_seed(seed=42):
    """设置随机种子，确保结果可复现"""
    random.seed(seed)
    np.random.seed(seed)

def create_split_directories(base_dir="animal_data"):
    """创建训练、验证和测试数据集目录"""
    # 创建基础目录
    if os.path.exists(base_dir):
        print(f"警告: {base_dir} 目录已存在，将被清空并重新创建")
        shutil.rmtree(base_dir)
    
    # 创建训练、验证和测试目录
    for split in ["train", "val", "test"]:
        os.makedirs(os.path.join(base_dir, split), exist_ok=True)
    
    return base_dir

def split_data(src_dir="data", dest_dir="animal_data", train_ratio=0.7, val_ratio=0.15, test_ratio=0.15, seed=42):
    """将数据集分割成训练、验证和测试集"""
    set_seed(seed)
    
    # 确保比例总和为1
    assert abs(train_ratio + val_ratio + test_ratio - 1.0) < 1e-5, "比例总和必须为1"
    
    # 创建目标目录
    dest_dir = create_split_directories(dest_dir)
    
    # 获取所有类别
    classes = [d for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))]
    print(f"发现 {len(classes)} 个类别: {classes}")
    
    # 为每个类别创建对应的子目录
    for split in ["train", "val", "test"]:
        for cls in classes:
            os.makedirs(os.path.join(dest_dir, split, cls), exist_ok=True)
    
    # 分割并复制数据
    for cls in classes:
        print(f"处理类别: {cls}")
        cls_dir = os.path.join(src_dir, cls)
        images = [img for img in os.listdir(cls_dir) if img.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        # 打乱图像顺序
        random.shuffle(images)
        
        # 计算每个集合的大小
        num_images = len(images)
        num_train = int(train_ratio * num_images)
        num_val = int(val_ratio * num_images)
        
        # 分割数据集
        train_images = images[:num_train]
        val_images = images[num_train:num_train+num_val]
        test_images = images[num_train+num_val:]
        
        # 复制图像到相应目录
        for subset, subset_images in [("train", train_images), ("val", val_images), ("test", test_images)]:
            print(f"  复制 {len(subset_images)} 张图像到 {subset} 集合")
            for img in tqdm(subset_images, desc=f"{cls}-{subset}"):
                src_path = os.path.join(cls_dir, img)
                dst_path = os.path.join(dest_dir, subset, cls, img)
                shutil.copy2(src_path, dst_path)
    
    # 打印统计信息
    print("\n数据集分割完成!")
    for split in ["train", "val", "test"]:
        split_classes = os.listdir(os.path.join(dest_dir, split))
        total_images = sum(len(os.listdir(os.path.join(dest_dir, split, cls))) for cls in split_classes)
        print(f"{split} 集合: {total_images} 张图像")

def main():
    # 设置参数
    src_dir = "data"  # 源数据目录
    dest_dir = "animal_data"  # 目标数据目录
    train_ratio = 0.7  # 训练集比例
    val_ratio = 0.15  # 验证集比例
    test_ratio = 0.15  # 测试集比例
    seed = 42  # 随机种子
    
    # 执行数据分割
    split_data(src_dir, dest_dir, train_ratio, val_ratio, test_ratio, seed)

if __name__ == "__main__":
    main()
